# import
import numpy as np
import pandas as pd

import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from util.models import *
from util.misc import *
from restrictiveness import *

## mat file
from scipy.io import loadmat

# load dataset
base = os.path.dirname(os.path.abspath(__file__))
file_name = os.path.normpath(os.path.join(base, '../data/Bruhin_et_al_2010.mat'))
matdata = loadmat(file_name)

df_full = pd.DataFrame()
for key in matdata.keys():
    if not '__' in key:
        df_full[key] = matdata[key].flatten()

## drop lotteries over losses
df = df_full[df_full['z2'] >= 0]

df_lottery = preprocess_df(df)
z1 = df_lottery['z1'].values
z2 = df_lottery['z2'].values
p1 = df_lottery['p1'].values
expval = df_lottery['expected_value'].values
base = os.path.dirname(os.path.abspath(__file__))
file_name = os.path.normpath(os.path.join(base, 'generated_ce.csv'))
df_ce = pd.read_csv(file_name)
ce_data = df_ce.values

# set search space for parameters
alpha_grid_size = 0.01
gamma_grid_size = 0.01
delta_grid_size = 0.1
alpha_ubd = 1
gamma_ubd = 1
delta_ubd = 30

alpha_grid = np.arange(alpha_grid_size, alpha_ubd+alpha_grid_size, alpha_grid_size)
gamma_grid = np.arange(gamma_grid_size, gamma_ubd+gamma_grid_size, gamma_grid_size)
delta_grid = np.array([1.0])
param_grid = cartesian_product(alpha_grid, gamma_grid, delta_grid)

restrictiveness, stderr = compute_restrictiveness(z1, z2, p1, expval, ce_data, param_grid)
print('restrictiveness: {}'.format(restrictiveness))
print('stderr: {}'.format(stderr))
